#include "my_pcl.h"

// 显式实例化 PCL 模板以解决链接问题
#include <pcl/impl/pcl_base.hpp>
#include <pcl/filters/impl/passthrough.hpp>
#include <pcl/filters/impl/voxel_grid.hpp>
#include <pcl/segmentation/impl/extract_clusters.hpp>
#include <pcl/filters/impl/extract_indices.hpp>
#include <pcl/registration/impl/icp.hpp>
#include <pcl/io/impl/pcd_io.hpp>

// 显式模板实例化
template class pcl::PCLBase<pcl::PointXYZI>;
template class pcl::PassThrough<pcl::PointXYZI>;
template class pcl::VoxelGrid<pcl::PointXYZI>;
template class pcl::EuclideanClusterExtraction<pcl::PointXYZI>;
template class pcl::ExtractIndices<pcl::PointXYZI>;
template class pcl::IterativeClosestPoint<pcl::PointXYZI, pcl::PointXYZI>;

// ==================== 全局发布器管理 ====================
std::map<std::string, ros::Publisher> global_publishers;
ros::NodeHandle* global_nh = nullptr;

// 获取或创建发布器
ros::Publisher& getOrCreatePublisher(const std::string& topic_name) {
    if (global_nh == nullptr) {
        // 创建默认NodeHandle
        static ros::NodeHandle default_nh;
        global_nh = &default_nh;
    }
    
    if (global_publishers.find(topic_name) == global_publishers.end()) {
        global_publishers[topic_name] = global_nh->advertise<sensor_msgs::PointCloud2>(topic_name, 1);
        ROS_INFO("Created publisher for topic: %s", topic_name.c_str());
    }
    
    return global_publishers[topic_name];
}

// ==================== 核心函数实现 ====================

// 类似imshow的便捷点云发布函数
void pclshow(const std::string& topic_name, 
             const pcl::PointCloud<PointType>::Ptr& cloud,
             uint8_t r, uint8_t g, uint8_t b,
             const std::string& frame_id) {
    
    if (cloud->empty()) {
        ROS_WARN("pclshow: Empty cloud for topic %s", topic_name.c_str());
        return;
    }
    
    // 获取发布器
    ros::Publisher& pub = getOrCreatePublisher(topic_name);
    
    // 创建header
    std_msgs::Header header;
    header.stamp = ros::Time::now();
    header.frame_id = frame_id;
    
    // 创建彩色点云
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr colored_cloud(new pcl::PointCloud<pcl::PointXYZRGB>);
    
    // 转换点云颜色
    for (const auto& point : cloud->points) {
        pcl::PointXYZRGB colored_point;
        colored_point.x = point.x;
        colored_point.y = point.y;
        colored_point.z = point.z;
        colored_point.r = r;
        colored_point.g = g;
        colored_point.b = b;
        colored_cloud->points.push_back(colored_point);
    }
    
    // 设置点云属性
    colored_cloud->width = colored_cloud->points.size();
    colored_cloud->height = 1;
    colored_cloud->is_dense = true;
    
    // 发布
    sensor_msgs::PointCloud2 cloud_msg;
    pcl::toROSMsg(*colored_cloud, cloud_msg);
    cloud_msg.header = header;
    pub.publish(cloud_msg);
    
    ROS_DEBUG("pclshow: Published %zu points to %s", cloud->points.size(), topic_name.c_str());
}

// 点云聚类检测函数
std::vector<pcl::PointCloud<PointType>::Ptr> processCloudsForStickDetection(
    pcl::PointCloud<PointType>::Ptr cloud,
    const std::string& publish_topic,
    const std::string& frame_id)
{
    // 1. 点云预处理 - 滤波和下采样
    pcl::PointCloud<PointType>::Ptr cloud_filtered(new pcl::PointCloud<PointType>);
    
    // 直通滤波 - X轴
    pcl::PassThrough<PointType> passx;
    passx.setInputCloud(cloud);
    passx.setFilterFieldName("x");
    passx.setFilterLimits(0.3, 4.0);
    passx.filter(*cloud_filtered);

    // 直通滤波 - Y轴
    pcl::PassThrough<PointType> passy;
    passy.setInputCloud(cloud_filtered);
    passy.setFilterFieldName("y");
    passy.setFilterLimits(-3.0, 0.0); 
    passy.filter(*cloud_filtered);

    // 体素网格下采样
    pcl::VoxelGrid<PointType> vg;
    vg.setInputCloud(cloud_filtered);
    vg.setLeafSize(0.003f, 0.003f, 0.003f);
    vg.filter(*cloud_filtered);
    
    // 2. 欧几里得聚类
    pcl::search::KdTree<PointType>::Ptr tree(new pcl::search::KdTree<PointType>());
    tree->setInputCloud(cloud_filtered);

    std::vector<pcl::PointIndices> initial_clusters;
    pcl::EuclideanClusterExtraction<PointType> ec;
    ec.setClusterTolerance(0.08);
    ec.setMinClusterSize(15);
    ec.setMaxClusterSize(10000);
    ec.setSearchMethod(tree);
    ec.setInputCloud(cloud_filtered);
    ec.extract(initial_clusters);
    
    // 限制聚类数量以提高实时性
    const size_t MAX_CLUSTERS = 30;
    if (initial_clusters.size() > MAX_CLUSTERS) {
        initial_clusters.resize(MAX_CLUSTERS);
    }

    // 3. 提取聚类并计算质心
    std::vector<pcl::PointCloud<PointType>::Ptr> cluster_clouds;
    std::vector<Eigen::Vector4f> cluster_centroids;
    std::vector<PointType> cluster_min_pts, cluster_max_pts;
    std::vector<bool> merged;
    
    // 预分配内存
    cluster_clouds.reserve(initial_clusters.size());
    cluster_centroids.reserve(initial_clusters.size());
    cluster_min_pts.reserve(initial_clusters.size());
    cluster_max_pts.reserve(initial_clusters.size());
    merged.reserve(initial_clusters.size());

    for (const auto& cluster_indices : initial_clusters)
    {
        pcl::PointCloud<PointType>::Ptr cluster(new pcl::PointCloud<PointType>);
        cluster->reserve(cluster_indices.indices.size());
        
        for (const auto& idx : cluster_indices.indices)
        {
            cluster->points.push_back(cloud_filtered->points[idx]);
        }
        cluster->width = cluster->points.size();
        cluster->height = 1;
        cluster->is_dense = true;

        // 计算质心和边界框
        Eigen::Vector4f centroid;
        PointType min_pt, max_pt;
        pcl::compute3DCentroid(*cluster, centroid);
        pcl::getMinMax3D(*cluster, min_pt, max_pt);
        
        cluster_clouds.push_back(cluster);
        cluster_centroids.push_back(centroid);
        cluster_min_pts.push_back(min_pt);
        cluster_max_pts.push_back(max_pt);
        merged.push_back(false);
    }

    // 4. 合并相邻聚类
    std::vector<pcl::PointCloud<PointType>::Ptr> merged_clusters;
    merged_clusters.reserve(initial_clusters.size()/2);

    const float vertical_merge_threshold = 0.15; // 15cm
    const float horizontal_merge_threshold = 0.08; // 8cm

    for (size_t i = 0; i < cluster_clouds.size(); i++)
    {
        if (merged[i])
            continue;

        pcl::PointCloud<PointType>::Ptr merged_cluster(new pcl::PointCloud<PointType>);
        *merged_cluster += *cluster_clouds[i];
        merged[i] = true;

        // 合并相邻聚类
        for (size_t j = i+1; j < cluster_clouds.size(); j++)
        {
            if (merged[j])
                continue;

            float x_dist = std::abs(cluster_centroids[i][0] - cluster_centroids[j][0]);
            float y_dist = std::abs(cluster_centroids[i][1] - cluster_centroids[j][1]);
            float horizontal_dist = std::sqrt(x_dist*x_dist + y_dist*y_dist);
            
            bool vertical_overlap = (cluster_min_pts[i].z <= cluster_max_pts[j].z + vertical_merge_threshold &&
                                     cluster_max_pts[i].z + vertical_merge_threshold >= cluster_min_pts[j].z);
            
            if (horizontal_dist < horizontal_merge_threshold && vertical_overlap)
            {
                *merged_cluster += *cluster_clouds[j];
                merged[j] = true;
                
                // 更新边界和质心
                pcl::getMinMax3D(*merged_cluster, cluster_min_pts[i], cluster_max_pts[i]);
                pcl::compute3DCentroid(*merged_cluster, cluster_centroids[i]);
            }
        }

        merged_clusters.push_back(merged_cluster);
    }
    
    // 5. 发布彩色聚类点云
    if (!publish_topic.empty() && !merged_clusters.empty()) {
        // 获取发布器
        ros::Publisher& pub = getOrCreatePublisher(publish_topic);
        
        // 创建header
        std_msgs::Header header;
        header.stamp = ros::Time::now();
        header.frame_id = frame_id;
        
        // 预定义颜色
        std::vector<std::array<uint8_t, 3>> colors = {
            {255, 0, 0},    {0, 255, 0},    {0, 0, 255},    {255, 255, 0},
            {255, 0, 255},  {0, 255, 255},  {255, 128, 0},  {128, 255, 0},
            {255, 128, 128},{128, 128, 255},{128, 255, 255},{255, 255, 128},
            {255, 128, 255},{128, 128, 128},{192, 192, 192}
        };
        
        // 合并所有聚类到彩色点云
        pcl::PointCloud<pcl::PointXYZRGB>::Ptr colored_cloud(new pcl::PointCloud<pcl::PointXYZRGB>);
        
        for (size_t i = 0; i < merged_clusters.size(); ++i) {
            const auto& cluster = merged_clusters[i];
            auto color = colors[i % colors.size()];
            
            for (const auto& point : cluster->points) {
                pcl::PointXYZRGB colored_point;
                colored_point.x = point.x;
                colored_point.y = point.y;
                colored_point.z = point.z;
                colored_point.r = color[0];
                colored_point.g = color[1];
                colored_point.b = color[2];
                colored_cloud->points.push_back(colored_point);
            }
        }
        
        // 设置点云属性
        colored_cloud->width = colored_cloud->points.size();
        colored_cloud->height = 1;
        colored_cloud->is_dense = true;
        
        // 发布
        sensor_msgs::PointCloud2 cloud_msg;
        pcl::toROSMsg(*colored_cloud, cloud_msg);
        cloud_msg.header = header;
        pub.publish(cloud_msg);
        
        ROS_DEBUG("Published %zu clusters with %zu total points to %s", 
                  merged_clusters.size(), colored_cloud->points.size(), publish_topic.c_str());
    }
    
    return merged_clusters;
}

// ==================== ICP_Matching类实现 ====================

ICP_Matching::ICP_Matching(const std::string& template_folder_path, double fitness_threshold) 
    : template_folder_path_(template_folder_path), icp_fitness_threshold_(fitness_threshold),
      max_iterations_(200), transformation_epsilon_(1e-5), euclidean_fitness_epsilon_(1e-5),
      max_correspondence_distance_(0.5), min_cluster_points_(50) {
    
    if (!template_folder_path_.empty() && template_folder_path_.back() != '/') {
        template_folder_path_ += '/';
    }
    
    ROS_INFO("ICP_Matching initialized with template folder: %s", template_folder_path_.c_str());
}

void ICP_Matching::setTemplateFiles(const std::vector<std::string>& template_files) {
    template_files_ = template_files;
    ROS_INFO("Loaded %zu template files for ICP matching", template_files_.size());
}

void ICP_Matching::setICPParameters(int max_iterations, double transformation_epsilon,
                                   double euclidean_fitness_epsilon, double max_correspondence_distance) {
    max_iterations_ = max_iterations;
    transformation_epsilon_ = transformation_epsilon;
    euclidean_fitness_epsilon_ = euclidean_fitness_epsilon;
    max_correspondence_distance_ = max_correspondence_distance;
    
    ROS_INFO("ICP parameters updated: max_iter=%d, trans_eps=%.2e, fit_eps=%.2e, max_corr_dist=%.3f",
             max_iterations_, transformation_epsilon_, euclidean_fitness_epsilon_, max_correspondence_distance_);
}

void ICP_Matching::setClusterFilterParameters(int min_cluster_points) {
    min_cluster_points_ = min_cluster_points;
    ROS_INFO("Cluster filter updated: min_points=%d", min_cluster_points_);
}

pcl::PointCloud<PointType>::Ptr ICP_Matching::loadTemplate(const std::string& template_file) {
    pcl::PointCloud<PointType>::Ptr template_cloud(new pcl::PointCloud<PointType>);
    
    std::string full_path = template_folder_path_ + template_file;
    
    if (pcl::io::loadPCDFile<PointType>(full_path, *template_cloud) == -1) {
        ROS_ERROR("Failed to load template file: %s", full_path.c_str());
        return nullptr;
    }
    
    ROS_DEBUG("Successfully loaded template: %s (%zu points)", template_file.c_str(), template_cloud->size());
    return template_cloud;
}

ICPMatchResult ICP_Matching::performICP(const pcl::PointCloud<PointType>::Ptr& template_cloud,
                                        const pcl::PointCloud<PointType>::Ptr& target_cluster) {
    ICPMatchResult result;
    
    // 配置ICP
    pcl::IterativeClosestPoint<PointType, PointType> icp;
    icp.setInputSource(template_cloud);
    icp.setInputTarget(target_cluster);
    
    // 设置ICP参数
    icp.setMaximumIterations(max_iterations_);
    icp.setTransformationEpsilon(transformation_epsilon_);
    icp.setEuclideanFitnessEpsilon(euclidean_fitness_epsilon_);
    icp.setMaxCorrespondenceDistance(max_correspondence_distance_);
    
    // 执行配准
    icp.align(*result.aligned_template);
    
    result.converged = icp.hasConverged();
    result.fitness_score = icp.getFitnessScore();
    result.transformation = icp.getFinalTransformation();
    
    // 计算配准后的位置
    if (result.converged) {
        Eigen::Vector4f centroid;
        pcl::compute3DCentroid(*result.aligned_template, centroid);
        result.x_position = centroid[0];
        result.y_position = centroid[1];
        result.z_position = centroid[2];
        
        // 计算物体属性
        calculateObjectProperties(result.aligned_template, result);
    }
    
    return result;
}

void ICP_Matching::calculateObjectProperties(const pcl::PointCloud<PointType>::Ptr& object_cloud, 
                                            ICPMatchResult& result) {
    // 使用PCA计算主方向
    Eigen::Vector4f centroid;
    pcl::compute3DCentroid(*object_cloud, centroid);
    
    // 计算协方差矩阵
    Eigen::Matrix3f covariance_matrix;
    pcl::computeCovarianceMatrixNormalized(*object_cloud, centroid, covariance_matrix);
    
    // 特征值分解
    Eigen::SelfAdjointEigenSolver<Eigen::Matrix3f> eigen_solver(covariance_matrix);
    Eigen::Matrix3f eigen_vectors = eigen_solver.eigenvectors();
    
    // 最小特征值对应的特征向量就是法向量
    result.normal_vector = eigen_vectors.col(0);
    
    // 确保法向量指向正方向
    if (result.normal_vector[2] < 0) {
        result.normal_vector = -result.normal_vector;
    }
    
    // 估算半径（基于点云分布）
    double max_distance = 0.0;
    for (const auto& point : object_cloud->points) {
        double distance = std::sqrt(
            std::pow(point.x - centroid[0], 2) + 
            std::pow(point.y - centroid[1], 2) + 
            std::pow(point.z - centroid[2], 2)
        );
        max_distance = std::max(max_distance, distance);
    }
    result.radius = max_distance * 0.8;  // 估算内径
}

TrajectoryPoints ICP_Matching::generateTrajectory(const ICPMatchResult& match_result, 
                                                  double approach_distance) {
    TrajectoryPoints trajectory;
    
    // 中心点
    trajectory.center_point.x = match_result.x_position;
    trajectory.center_point.y = match_result.y_position;
    trajectory.center_point.z = match_result.z_position;
    
    // 法向量方向的接近点
    trajectory.approach_point.x = match_result.x_position - match_result.normal_vector[0] * approach_distance;
    trajectory.approach_point.y = match_result.y_position - match_result.normal_vector[1] * approach_distance;
    trajectory.approach_point.z = match_result.z_position - match_result.normal_vector[2] * approach_distance;
    
    // 法向量反方向的离开点
    trajectory.exit_point.x = match_result.x_position + match_result.normal_vector[0] * approach_distance;
    trajectory.exit_point.y = match_result.y_position + match_result.normal_vector[1] * approach_distance;
    trajectory.exit_point.z = match_result.z_position + match_result.normal_vector[2] * approach_distance;
    
    // 计算总距离
    trajectory.total_distance = 2.0 * approach_distance;
    
    return trajectory;
}

std::vector<ICPMatchResult> ICP_Matching::findMatches(
    const std::vector<pcl::PointCloud<PointType>::Ptr>& clusters,
    bool visualize_templates) {
    
    std::vector<ICPMatchResult> all_results;
    
    ROS_INFO("========== ICP模板匹配开始 ==========");
    ROS_INFO("总聚类数: %zu, 模板数: %zu", clusters.size(), template_files_.size());
    
    int cluster_index = 0;
    for (const auto& cluster : clusters) {
        // 过滤点云数过少的聚类
        if (cluster->size() < min_cluster_points_) {
            ROS_DEBUG("跳过点云数过少的聚类 %d (点数: %zu)", cluster_index, cluster->size());
            cluster_index++;
            continue;
        }
        
        ROS_INFO("处理聚类 %d (点数: %zu)", cluster_index, cluster->size());

        ICPMatchResult best_match;
        
        // 尝试不同的模板文件
        for (size_t template_idx = 0; template_idx < template_files_.size(); ++template_idx) {
            const std::string& template_file = template_files_[template_idx];
            
            // 加载模板
            auto template_cloud = loadTemplate(template_file);
            if (template_cloud == nullptr) {
                ROS_WARN("跳过无效模板文件: %s", template_file.c_str());
                continue;
            }
            
            // 可视化模板
            if (visualize_templates) {
                std::string template_topic = "/template_" + std::to_string(template_idx) + "_" + 
                                            template_file.substr(0, template_file.find(".pcd"));
                pclshow(template_topic, template_cloud, 0, 255, 255); // 青色显示模板
            }
            
            // 执行ICP匹配
            ICPMatchResult current_match = performICP(template_cloud, cluster);
            
            if (current_match.converged) {
                ROS_INFO("  模板 (%s) - 匹配度: %.6f, 收敛: %s", 
                         template_file.c_str(), current_match.fitness_score,
                         current_match.converged ? "是" : "否");
            }
            
            // 保存最佳匹配结果
            if (current_match.converged && current_match.fitness_score < best_match.fitness_score) {
                best_match = current_match;
            }
        }
        
        // 如果找到好的匹配结果，添加到结果列表
        if (best_match.converged && best_match.fitness_score < icp_fitness_threshold_) {
            all_results.push_back(best_match);
            ROS_INFO("  ✓ 找到物体匹配! 位置: (%.3f, %.3f, %.3f), 匹配度: %.6f", 
                     best_match.x_position, best_match.y_position, best_match.z_position,
                     best_match.fitness_score);
        }
        
        cluster_index++;
    }
    
    ROS_INFO("========== ICP模板匹配结束 ==========");
    ROS_INFO("找到的物体数量: %zu", all_results.size());
    
    return all_results;
}

void ICP_Matching::publishVisualization(const std::vector<ICPMatchResult>& matches,
                                        const std_msgs::Header& header,
                                        const std::string& detected_topic,
                                        const std::string& trajectory_topic) {
    if (matches.empty()) return;
    
    // 获取发布器
    ros::Publisher& detected_pub = getOrCreatePublisher(detected_topic);
    ros::Publisher& trajectory_pub = getOrCreatePublisher(trajectory_topic);
    
    // 预定义颜色
    std::vector<std::array<uint8_t, 3>> object_colors = {
        {255, 0, 0}, {0, 255, 0}, {0, 0, 255}, {255, 255, 0}
    };
    
    std::vector<std::array<uint8_t, 3>> trajectory_colors = {
        {255, 0, 0},    // 接近点 - 红色
        {0, 255, 0},    // 中心点 - 绿色
        {0, 0, 255}     // 离开点 - 蓝色
    };
    
    // 创建检测结果点云
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr detected_cloud(new pcl::PointCloud<pcl::PointXYZRGB>);
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr trajectory_cloud(new pcl::PointCloud<pcl::PointXYZRGB>);
    
    for (size_t i = 0; i < matches.size(); ++i) {
        const auto& match = matches[i];
        auto color = object_colors[i % object_colors.size()];
        
        // 添加检测结果点云
        for (const auto& point : match.aligned_template->points) {
            pcl::PointXYZRGB colored_point;
            colored_point.x = point.x;
            colored_point.y = point.y;
            colored_point.z = point.z;
            colored_point.r = color[0];
            colored_point.g = color[1];
            colored_point.b = color[2];
            detected_cloud->points.push_back(colored_point);
        }
        
        // 生成并添加轨迹点
        TrajectoryPoints trajectory = generateTrajectory(match, 0.5);
        
        // 接近点
        pcl::PointXYZRGB approach_point;
        approach_point.x = trajectory.approach_point.x;
        approach_point.y = trajectory.approach_point.y;
        approach_point.z = trajectory.approach_point.z;
        approach_point.r = trajectory_colors[0][0];
        approach_point.g = trajectory_colors[0][1];
        approach_point.b = trajectory_colors[0][2];
        trajectory_cloud->points.push_back(approach_point);
        
        // 中心点
        pcl::PointXYZRGB center_point;
        center_point.x = trajectory.center_point.x;
        center_point.y = trajectory.center_point.y;
        center_point.z = trajectory.center_point.z;
        center_point.r = trajectory_colors[1][0];
        center_point.g = trajectory_colors[1][1];
        center_point.b = trajectory_colors[1][2];
        trajectory_cloud->points.push_back(center_point);
        
        // 离开点
        pcl::PointXYZRGB exit_point;
        exit_point.x = trajectory.exit_point.x;
        exit_point.y = trajectory.exit_point.y;
        exit_point.z = trajectory.exit_point.z;
        exit_point.r = trajectory_colors[2][0];
        exit_point.g = trajectory_colors[2][1];
        exit_point.b = trajectory_colors[2][2];
        trajectory_cloud->points.push_back(exit_point);
        
        // 单独发布每个物体模板以便可视化
        std::string individual_topic = "/object_template_" + std::to_string(i);
        pclshow(individual_topic, match.aligned_template, color[0], color[1], color[2], header.frame_id);
    }
    
    // 设置点云属性并发布
    detected_cloud->width = detected_cloud->points.size();
    detected_cloud->height = 1;
    detected_cloud->is_dense = true;
    
    trajectory_cloud->width = trajectory_cloud->points.size();
    trajectory_cloud->height = 1;
    trajectory_cloud->is_dense = true;
    
    // 发布检测结果
    if (!detected_cloud->points.empty()) {
        sensor_msgs::PointCloud2 detected_msg;
        pcl::toROSMsg(*detected_cloud, detected_msg);
        detected_msg.header = header;
        detected_pub.publish(detected_msg);
    }
    
    // 发布轨迹
    if (!trajectory_cloud->points.empty()) {
        sensor_msgs::PointCloud2 trajectory_msg;
        pcl::toROSMsg(*trajectory_cloud, trajectory_msg);
        trajectory_msg.header = header;
        trajectory_pub.publish(trajectory_msg);
    }
    
    ROS_INFO("已发布 %zu 个物体的检测结果和轨迹", matches.size());
}
